#include "pullFromRtsp.h"

namespace sp::Image::Pull{

void register_pullFromRtsp(){
    pullFromRtsp::registe();
}

pullFromRtsp::pullFromRtsp(std::vector<std::string> v){
    assert(v.size() == 4);
    name = "pullFromRtsp";
    class_name = "Image-Pull-pullFromRtsp";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({},MAT));
    this->rtspUrl = v[0];
    this->width = std::stoi(v[1]);
    this->height = std::stoi(v[2]);
    this->size = std::stol(v[3]);
    this->index = 0;
}

pullFromRtsp::pullFromRtsp(std::string url, int width, int height, long size){
    name = "pullFromRtsp";
    class_name = "Image-Pull-pullFromRtsp";
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({},MAT));
    this->rtspUrl = url;
    this->width = width;
    this->height = height;
    this->size = size;
    this->index = 0;
}

pullFromRtsp::pullFromRtsp(){
    name = "pullFromRtsp";
    class_name = "Image-Pull-pullFromRtsp";
    this->index = 0;
    data_type_support.push_back(std::pair<std::vector<uint8_t>, uint8_t>({},MAT));
}

pullFromRtsp::~pullFromRtsp(){

}

void pullFromRtsp::registe(){
    Function::Register("Image-Pull-pullFromRtsp", createObject);
    Function::Register("Image-Pull-pullFromRtsp", createObject2);
}

Function* pullFromRtsp::createObject(std::vector<std::string> params){
    return new pullFromRtsp(params);
}

Function* pullFromRtsp::createObject2(){
    return new pullFromRtsp();
}

void pullFromRtsp::start(){
    capture = cv::VideoCapture(rtspUrl);
    if(!capture.isOpened()){
        std::cout << "video open failed!" << std::endl;
    }
}

bool pullFromRtsp::waitForResource(){
    return true;
}

void pullFromRtsp::process(std::vector<std::vector<Data*>>& data_input, std::vector<Data*>& data_output){
    // 验证输入个数
    assert(batch_size != 0);
    assert(data_input.size() <= batch_size);
    assert(data_input[0].size() == 1);
    // 处理
    for(int i = 0; i < data_input.size(); i++){
        if(index >= size){
            // 释放内存
            executor->free((char*)data_input[i][0]);
            data_output.push_back(DATA_END);
            break;
        }else{
            // 从Rtsp流中获取一帧
            char* mat_data_ptr = executor->malloc(width, height);
            cv::Mat* frame = new cv::Mat(height, width, CV_8UC3);
            frame->data = (uchar*)mat_data_ptr;
            frame->datastart = frame->data;
            frame->dataend = frame->datastart + height*width*3;
            frame->datalimit = frame->dataend;
            capture >> *frame;
            if(frame->empty()){
                data_output.push_back(DATA_END);
                break;
            }else{
                data_input[i][0]->set(3, height, width, (uchar*)(frame->data));
                data_output.push_back(data_input[i][0]);
                // std::cout << "pull frame " << width << "*" << height << " : " << index << std::endl;
            }
        }
        index++;
    }
}

bool pullFromRtsp::releaseResource(){
    return true;
}

void pullFromRtsp::finish(){
   capture.release();
}

void pullFromRtsp::copy(Function* other){
    Function::defaultCopy(other);
    pullFromRtsp* oth = dynamic_cast<pullFromRtsp*>(other);
    this->rtspUrl = oth->rtspUrl;
    this->width = oth->width;
    this->height = oth->height;
    this->size = oth->size;
}

};